import cv2
import mediapipe as mp
import time

from HandTrakingMin import imgRGB


class HandDetector():
    def __init__(self, mode=False, HandsMax=2, detectionCon=0.5, trackCon=0.5):
        """
        初始化参数
        :param mode: 是否输入静态图像
        :param HandsMax: 检测到手的最大数量
        :param detectionCon: 检测手的置信度
        :param trackCon: 追踪手的置信度
        """
        self.results = self.hands.process(imgRGB)
        self.mode = mode
        self.HandMax = HandsMax
        self.detectionCon = detectionCon
        self.trackCon = trackCon

        self.mpHands = mp.solutions.hands
        self.hands = self.mpHands.Hands(self.mode, self.HandMax,
                                        self.detectionCon, self.trackCon)
        self.mpDraw = mp.solutions.drawing_utils

    def findHands(self, img1, draw=True):
        """
        检测手掌
        :param img1:要识别的一帧图像
        :param draw:是否对手的标志点进行绘图
        :return:绘画完成的一帧图像
        """
        imgRGB = cv2.cvtColor(img1, cv2.COLOR_BGR2RGB)
        if self.results.multi_hand_landmarks:
            for handLms in self.results.multi_hand_landmarks:
                if draw:
                    self.mpDraw.draw_landmarks(img1, handLms, self.mpHands.HAND_CONNECTIONS)
        return img1

    def findPosition(self, img1, handNo=0, personDraw=True):
        """
        检测手的标志点
        :param img1: 要识别的一帧图像
        :param handNo: 手的编号
        :param personDraw: 是否对手的标志点进行个性化绘图
        :return: 手的21个标志点位置
        """
        lmList1 = []
        if self.results.multi_hand_landmarks:
            MyHand = self.results.multi_hand_landmarks[handNo]
            for id1, lm in enumerate(MyHand.landmark):
                h, w, c = img1.shape
                cx, cy = int(lm.x * w), int(lm.y * h)
                lmList1.append([id1, cx, cy])
                if personDraw:
                    cv2.circle(img1, (cx, cy), int(w / 50), (255, 0, 255), cv2.FILLED)
        return lmList1


if __name__ == "__main__":

    pTime = 0  # 开始时间初始化

    cap = cv2.VideoCapture("example.mp4")
    detector = HandDetector()

    while True:
        success, img = cap.read()
        img = detector.findHands(img, draw=True)
        lmList = detector.findPosition(img, personDraw=False)
        if len(lmList) != 0:
            print(lmList[4])  # 食指最上部点的坐标

        cTime = time.time()
        fps = 1 / (cTime - pTime)
        pTime = cTime

        cv2.putText(img, str(int(fps)), (10, 70), cv2.FONT_HERSHEY_PLAIN, 3,
                    (255, 0, 255), 1)
        cv2.imshow("Image", img)
        cv2.waitKey(1)


